#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Time    : 2021/7/27 1:23 上午
# @Author  : WangZhixing
import torch
import torch.nn.functional as F

from torch_geometric.nn import GATConv


class GAT(torch.nn.Module):
    def __init__(self, num_features, hidden_channels, num_classes, n_heads=1):
        super(GAT, self).__init__()
        torch.manual_seed(12345)
        self.conv1 = GATConv(num_features, hidden_channels, heads=n_heads)
        self.conv2 = GATConv(n_heads * hidden_channels, num_classes)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = x.relu()
        #激活函数，舍弃所有0值
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.conv2(x, edge_index)
        return x
